{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "在本指南中，我们将探讨使用create_sql_query_chain来改善SQL查询生成的提示策略。我们将重点介绍如何在提示中获取与数据库相关的信息。\n",
    "\n",
    "我们将涵盖以下内容：\n",
    "\n",
    "LangChain SQLDatabase的方言如何影响提示链； 如何使用SQLDatabase.get_context将模式信息格式化到提示中； 如何构建和选择少量示例以协助模型。"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fd530944f37c0996"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sqlite\n",
      "['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']\n",
      "[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "db = SQLDatabase.from_uri(os.getenv(\"SQLlITE_DATABASE_URL\"), sample_rows_in_table_info=3)\n",
    "\n",
    "\n",
    "def get_schema(_):\n",
    "    return db.get_table_info()\n",
    "\n",
    "\n",
    "def run_query(query):\n",
    "    return db.run(query)\n",
    "\n",
    "\n",
    "print(db.dialect)\n",
    "print(db.get_usable_table_names())\n",
    "print(db.run(\"SELECT * FROM Artist LIMIT 10;\"))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:48:45.521535Z",
     "start_time": "2024-11-01T02:48:45.480103Z"
    }
   },
   "id": "64bcf604bab7eae9",
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Dialect-specific prompting\n",
    "最简单的事情之一是让我们的提示针对我们正在使用的SQL方言。当使用内置的create_sql_query_chain和SQLDatabase时，无论使用哪种方言，这一点都会得到处理：\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aec71d8281269841"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "['crate',\n 'duckdb',\n 'googlesql',\n 'mssql',\n 'mysql',\n 'mariadb',\n 'oracle',\n 'postgresql',\n 'sqlite',\n 'clickhouse',\n 'prestodb']"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains.sql_database.prompt import SQL_PROMPTS\n",
    "\n",
    "list(SQL_PROMPTS)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:50:02.719779Z",
     "start_time": "2024-11-01T02:50:01.831227Z"
    }
   },
   "id": "9cdbbea046047f4c",
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, first create a syntactically correct SQLite query to run, then look at the results of the query and return the answer to the input question.\n",
      "Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the LIMIT clause as per SQLite. You can order the results to return the most informative data in the database.\n",
      "Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in double quotes (\") to denote them as delimited identifiers.\n",
      "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
      "Pay attention to use date('now') function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: Question here\n",
      "SQLQuery: SQL Query to run\n",
      "SQLResult: Result of the SQLQuery\n",
      "Answer: Final answer here\n",
      "\n",
      "Only use the following tables:\n",
      "\u001B[33;1m\u001B[1;3m{table_info}\u001B[0m\n",
      "\n",
      "Question: \u001B[33;1m\u001B[1;3m{input}\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain.chains.sql_database.query import create_sql_query_chain\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\",\n",
    "    verbose=True\n",
    ")\n",
    "chain = create_sql_query_chain(llm, db)\n",
    "chain.get_prompts()[0].pretty_print()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:51:53.533026Z",
     "start_time": "2024-11-01T02:51:52.673790Z"
    }
   },
   "id": "c17fa5c1ba02761a",
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "**API Reference:**[create_sql_query_chain](https://python.langchain.com/api_reference/langchain/chains/langchain.chains.sql_database.query.create_sql_query_chain.html)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9f4d1ce9dc2ad23a"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 表定义和示例行\n",
    "获取DB上下文，查看表信息`SQLDatabase.get_context`\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b8c6a073f53533e8"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['table_info', 'table_names']\n",
      "\n",
      "CREATE TABLE \"Album\" (\n",
      "\t\"AlbumId\" INTEGER NOT NULL, \n",
      "\t\"Title\" NVARCHAR(160) NOT NULL, \n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"AlbumId\"), \n",
      "\tFOREIGN KEY(\"ArtistId\") REFERENCES \"Artist\" (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Album table:\n",
      "AlbumId\tTitle\tArtistId\n",
      "1\tFor Those About To Rock We Salute You\t1\n",
      "2\tBalls to the Wall\t2\n",
      "3\tRestless and Wild\t2\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Artist\" (\n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Artist table:\n",
      "ArtistId\tName\n",
      "1\tAC/DC\n",
      "2\tAccept\n",
      "3\tAerosmith\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Customer\" (\n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"FirstName\" NVARCHAR(40) NOT NULL, \n",
      "\t\"LastName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"Company\" NVARCHAR(80), \n",
      "\t\"Address\" NVARCHAR(70), \n",
      "\t\"City\" NVARCHAR(40), \n",
      "\t\"State\" NVARCHAR(40), \n",
      "\t\"Country\" NVARCHAR(40), \n",
      "\t\"PostalCode\" NVARCHAR(10), \n",
      "\t\"Phone\" NVARCHAR(24), \n",
      "\t\"Fax\" NVARCHAR(24), \n",
      "\t\"Email\" NVARCHAR(60) NOT NULL, \n",
      "\t\"SupportRepId\" INTEGER, \n",
      "\tPRIMARY KEY (\"CustomerId\"), \n",
      "\tFOREIGN KEY(\"SupportRepId\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Customer table:\n",
      "CustomerId\tFirstName\tLastName\tCompany\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\tSupportRepId\n",
      "1\tLuís\tGonçalves\tEmbraer - Empresa Brasileira de Aeronáutica S.A.\tAv. Brigadeiro Faria Lima, 2170\tSão José dos Campos\tSP\tBrazil\t12227-000\t+55 (12) 3923-5555\t+55 (12) 3923-5566\tluisg@embraer.com.br\t3\n",
      "2\tLeonie\tKöhler\tNone\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t+49 0711 2842222\tNone\tleonekohler@surfeu.de\t5\n",
      "3\tFrançois\tTremblay\tNone\t1498 rue Bélanger\tMontréal\tQC\tCanada\tH2G 1A7\t+1 (514) 721-4711\tNone\tftremblay@gmail.com\t3\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Employee\" (\n",
      "\t\"EmployeeId\" INTEGER NOT NULL, \n",
      "\t\"LastName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"FirstName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"Title\" NVARCHAR(30), \n",
      "\t\"ReportsTo\" INTEGER, \n",
      "\t\"BirthDate\" DATETIME, \n",
      "\t\"HireDate\" DATETIME, \n",
      "\t\"Address\" NVARCHAR(70), \n",
      "\t\"City\" NVARCHAR(40), \n",
      "\t\"State\" NVARCHAR(40), \n",
      "\t\"Country\" NVARCHAR(40), \n",
      "\t\"PostalCode\" NVARCHAR(10), \n",
      "\t\"Phone\" NVARCHAR(24), \n",
      "\t\"Fax\" NVARCHAR(24), \n",
      "\t\"Email\" NVARCHAR(60), \n",
      "\tPRIMARY KEY (\"EmployeeId\"), \n",
      "\tFOREIGN KEY(\"ReportsTo\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Employee table:\n",
      "EmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n",
      "1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n",
      "2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n",
      "3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Genre\" (\n",
      "\t\"GenreId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"GenreId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Genre table:\n",
      "GenreId\tName\n",
      "1\tRock\n",
      "2\tJazz\n",
      "3\tMetal\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Invoice\" (\n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceDate\" DATETIME NOT NULL, \n",
      "\t\"BillingAddress\" NVARCHAR(70), \n",
      "\t\"BillingCity\" NVARCHAR(40), \n",
      "\t\"BillingState\" NVARCHAR(40), \n",
      "\t\"BillingCountry\" NVARCHAR(40), \n",
      "\t\"BillingPostalCode\" NVARCHAR(10), \n",
      "\t\"Total\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceId\"), \n",
      "\tFOREIGN KEY(\"CustomerId\") REFERENCES \"Customer\" (\"CustomerId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Invoice table:\n",
      "InvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n",
      "1\t2\t2021-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n",
      "2\t4\t2021-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n",
      "3\t8\t2021-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"InvoiceLine\" (\n",
      "\t\"InvoiceLineId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\t\"Quantity\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceLineId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"InvoiceId\") REFERENCES \"Invoice\" (\"InvoiceId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from InvoiceLine table:\n",
      "InvoiceLineId\tInvoiceId\tTrackId\tUnitPrice\tQuantity\n",
      "1\t1\t2\t0.99\t1\n",
      "2\t1\t4\t0.99\t1\n",
      "3\t2\t6\t0.99\t1\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"MediaType\" (\n",
      "\t\"MediaTypeId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"MediaTypeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from MediaType table:\n",
      "MediaTypeId\tName\n",
      "1\tMPEG audio file\n",
      "2\tProtected AAC audio file\n",
      "3\tProtected MPEG-4 video file\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Playlist\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Playlist table:\n",
      "PlaylistId\tName\n",
      "1\tMusic\n",
      "2\tMovies\n",
      "3\tTV Shows\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from PlaylistTrack table:\n",
      "PlaylistId\tTrackId\n",
      "1\t3402\n",
      "1\t3389\n",
      "1\t3390\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Track\" (\n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(200) NOT NULL, \n",
      "\t\"AlbumId\" INTEGER, \n",
      "\t\"MediaTypeId\" INTEGER NOT NULL, \n",
      "\t\"GenreId\" INTEGER, \n",
      "\t\"Composer\" NVARCHAR(220), \n",
      "\t\"Milliseconds\" INTEGER NOT NULL, \n",
      "\t\"Bytes\" INTEGER, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"MediaTypeId\") REFERENCES \"MediaType\" (\"MediaTypeId\"), \n",
      "\tFOREIGN KEY(\"GenreId\") REFERENCES \"Genre\" (\"GenreId\"), \n",
      "\tFOREIGN KEY(\"AlbumId\") REFERENCES \"Album\" (\"AlbumId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Track table:\n",
      "TrackId\tName\tAlbumId\tMediaTypeId\tGenreId\tComposer\tMilliseconds\tBytes\tUnitPrice\n",
      "1\tFor Those About To Rock (We Salute You)\t1\t1\t1\tAngus Young, Malcolm Young, Brian Johnson\t343719\t11170334\t0.99\n",
      "2\tBalls to the Wall\t2\t2\t1\tU. Dirkschneider, W. Hoffmann, H. Frank, P. Baltes, S. Kaufmann, G. Hoffmann\t342562\t5510424\t0.99\n",
      "3\tFast As a Shark\t3\t2\t1\tF. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman\t230619\t3990994\t0.99\n",
      "*/\n"
     ]
    }
   ],
   "source": [
    "context = db.get_context()\n",
    "print(list(context))\n",
    "print(context[\"table_info\"])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:52:35.374033Z",
     "start_time": "2024-11-01T02:52:35.336294Z"
    }
   },
   "id": "e3b2317a79381bec",
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "当表格不多或表格宽度不大时，我们可以将所有这些信息插入到提示中："
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a886bf8229e77f49"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, first create a syntactically correct SQLite query to run, then look at the results of the query and return the answer to the input question.\n",
      "Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the LIMIT clause as per SQLite. You can order the results to return the most informative data in the database.\n",
      "Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in double quotes (\") to denote them as delimited identifiers.\n",
      "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
      "Pay attention to use date('now') function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: Question here\n",
      "SQLQuery: SQL Query to run\n",
      "SQLResult: Result of the SQLQuery\n",
      "Answer: Final answer here\n",
      "\n",
      "Only use the following tables:\n",
      "\n",
      "CREATE TABLE \"Album\" (\n",
      "\t\"AlbumId\" INTEGER NOT NULL, \n",
      "\t\"Title\" NVARCHAR(160) NOT NULL, \n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"AlbumId\"), \n",
      "\tFOREIGN KEY(\"ArtistId\") REFERENCES \"Artist\" (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Album table:\n",
      "AlbumId\tTitle\tArtistId\n",
      "1\tFor Those About To Rock We Salute You\t1\n",
      "2\tBalls to the Wall\t2\n",
      "3\tRestless and Wild\t2\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Artist\" (\n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120)\n"
     ]
    }
   ],
   "source": [
    "prompt_with_context = chain.get_prompts()[0].partial(table_info=context[\"table_info\"])\n",
    "print(prompt_with_context.pretty_repr()[:1500])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:54:07.182539Z",
     "start_time": "2024-11-01T02:54:07.169915Z"
    }
   },
   "id": "c0dc48219fc6544e",
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 样例"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c198b05320f70c11"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\"input\": \"List all artists.\", \"query\": \"SELECT * FROM Artist;\"},\n",
    "    {\n",
    "        \"input\": \"Find all albums for the artist 'AC/DC'.\",\n",
    "        \"query\": \"SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all tracks in the 'Rock' genre.\",\n",
    "        \"query\": \"SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find the total duration of all tracks.\",\n",
    "        \"query\": \"SELECT SUM(Milliseconds) FROM Track;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all customers from Canada.\",\n",
    "        \"query\": \"SELECT * FROM Customer WHERE Country = 'Canada';\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many tracks are there in the album with ID 5?\",\n",
    "        \"query\": \"SELECT COUNT(*) FROM Track WHERE AlbumId = 5;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find the total number of invoices.\",\n",
    "        \"query\": \"SELECT COUNT(*) FROM Invoice;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all tracks that are longer than 5 minutes.\",\n",
    "        \"query\": \"SELECT * FROM Track WHERE Milliseconds > 300000;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Who are the top 5 customers by total purchase?\",\n",
    "        \"query\": \"SELECT CustomerId, SUM(Total) AS TotalPurchase FROM Invoice GROUP BY CustomerId ORDER BY TotalPurchase DESC LIMIT 5;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Which albums are from the year 2000?\",\n",
    "        \"query\": \"SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many employees are there\",\n",
    "        \"query\": 'SELECT COUNT(*) FROM \"Employee\"',\n",
    "    },\n",
    "]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:55:00.400018Z",
     "start_time": "2024-11-01T02:55:00.394738Z"
    }
   },
   "id": "7aab82151018fdcd",
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "prompt = FewShotPromptTemplate(\n",
    "    examples=examples[:5],\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than {top_k} rows.\\n\\nHere is the relevant table info: {table_info}\\n\\nBelow are a number of examples of questions and their corresponding SQL queries.\",\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:55:57.628935Z",
     "start_time": "2024-11-01T02:55:57.621396Z"
    }
   },
   "id": "d02d555aa662a67b",
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "source": [
    "**API Reference:**[FewShotPromptTemplate](https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.few_shot.FewShotPromptTemplate.html) | [PromptTemplate](https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.prompt.PromptTemplate.html)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "cb62f5067bb69bca"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than 3 rows.\n",
      "\n",
      "Here is the relevant table info: foo\n",
      "\n",
      "Below are a number of examples of questions and their corresponding SQL queries.\n",
      "\n",
      "User input: List all artists.\n",
      "SQL query: SELECT * FROM Artist;\n",
      "\n",
      "User input: Find all albums for the artist 'AC/DC'.\n",
      "SQL query: SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');\n",
      "\n",
      "User input: List all tracks in the 'Rock' genre.\n",
      "SQL query: SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\n",
      "\n",
      "User input: Find the total duration of all tracks.\n",
      "SQL query: SELECT SUM(Milliseconds) FROM Track;\n",
      "\n",
      "User input: List all customers from Canada.\n",
      "SQL query: SELECT * FROM Customer WHERE Country = 'Canada';\n",
      "\n",
      "User input: How many artists are there?\n",
      "SQL query: \n"
     ]
    }
   ],
   "source": [
    "print(prompt.format(input=\"How many artists are there?\", top_k=3, table_info=\"foo\"))\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T02:57:02.794550Z",
     "start_time": "2024-11-01T02:57:02.790727Z"
    }
   },
   "id": "8b836bc58fda916a",
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Dynamic few-shot examples\n",
    "可以使用ExampleSelector来完成这个工作。在这种情况下，我们将使用SemanticSimilarityExampleSelector，它将例子保存在我们选择的向量数据库中。在运行时，它将执行输入与例子之间的相似度搜索，并返回语义上最相似的例子。\n",
    "\n",
    "在这里，我们默认使用OpenAI嵌入，但您可以将其替换为您选择的模型提供者的嵌入。"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "27a9824959fbfb08"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from langchain_community.embeddings import DashScopeEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.example_selectors import SemanticSimilarityExampleSelector\n",
    "\n",
    "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    examples,\n",
    "    DashScopeEmbeddings(dashscope_api_key=os.getenv(\"DASHSCOPE_API_KEY\")),\n",
    "    FAISS,\n",
    "    k=5,\n",
    "    input_keys=[\"input\"],\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T03:00:06.727400Z",
     "start_time": "2024-11-01T03:00:03.950711Z"
    }
   },
   "id": "15a5bed72633986a",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "[{'input': 'List all artists.', 'query': 'SELECT * FROM Artist;'},\n {'input': 'How many tracks are there in the album with ID 5?',\n  'query': 'SELECT COUNT(*) FROM Track WHERE AlbumId = 5;'},\n {'input': 'How many employees are there',\n  'query': 'SELECT COUNT(*) FROM \"Employee\"'},\n {'input': 'List all tracks that are longer than 5 minutes.',\n  'query': 'SELECT * FROM Track WHERE Milliseconds > 300000;'},\n {'input': \"List all tracks in the 'Rock' genre.\",\n  'query': \"SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\"}]"
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_selector.select_examples({\"input\": \"how many artists are there?\"})"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T03:00:49.797825Z",
     "start_time": "2024-11-01T03:00:49.170864Z"
    }
   },
   "id": "2d19e705e4e21d4e",
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "source": [
    "要使用它，我们可以直接将ExampleSelector传入我们的FewShotPromptTemplate中：\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "51155a3e6c85904d"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than {top_k} rows.\\n\\nHere is the relevant table info: {table_info}\\n\\nBelow are a number of examples of questions and their corresponding SQL queries.\",\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T03:01:59.898263Z",
     "start_time": "2024-11-01T03:01:59.893446Z"
    }
   },
   "id": "78796a2e8a03ce46",
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than 3 rows.\n",
      "\n",
      "Here is the relevant table info: foo\n",
      "\n",
      "Below are a number of examples of questions and their corresponding SQL queries.\n",
      "\n",
      "User input: List all artists.\n",
      "SQL query: SELECT * FROM Artist;\n",
      "\n",
      "User input: How many tracks are there in the album with ID 5?\n",
      "SQL query: SELECT COUNT(*) FROM Track WHERE AlbumId = 5;\n",
      "\n",
      "User input: How many employees are there\n",
      "SQL query: SELECT COUNT(*) FROM \"Employee\"\n",
      "\n",
      "User input: List all tracks that are longer than 5 minutes.\n",
      "SQL query: SELECT * FROM Track WHERE Milliseconds > 300000;\n",
      "\n",
      "User input: List all tracks in the 'Rock' genre.\n",
      "SQL query: SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\n",
      "\n",
      "User input: how many artists are there?\n",
      "SQL query: \n"
     ]
    }
   ],
   "source": [
    "print(prompt.format(input=\"how many artists are there?\", top_k=3, table_info=\"foo\"))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T03:02:25.877013Z",
     "start_time": "2024-11-01T03:02:25.499096Z"
    }
   },
   "id": "5800c758dd6eb076",
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "'SQL query: SELECT COUNT(*) FROM Artist;'"
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = create_sql_query_chain(llm, db, prompt)\n",
    "chain.invoke({\"question\": \"how many artists are there?\"})"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-01T03:02:49.625176Z",
     "start_time": "2024-11-01T03:02:46.240700Z"
    }
   },
   "id": "4cd587acd45fca3d",
   "execution_count": 14
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
